package cn.ddiancan.gateway.config;

import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import cn.ddiancan.xddcloud.common.entity.CommonResult;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayFlowRule;
import com.alibaba.csp.sentinel.adapter.gateway.common.rule.GatewayRuleManager;
import com.alibaba.csp.sentinel.adapter.gateway.sc.SentinelGatewayFilter;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.BlockRequestHandler;
import com.alibaba.csp.sentinel.adapter.gateway.sc.callback.GatewayCallbackManager;
import com.alibaba.csp.sentinel.adapter.gateway.sc.exception.SentinelGatewayBlockExceptionHandler;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.client.discovery.ReactiveDiscoveryClient;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.ServerResponse;
import org.springframework.web.reactive.result.view.ViewResolver;
import reactor.core.publisher.Flux;

@Configuration
@Slf4j
public class GatewaySentinelConfiguration implements ApplicationContextAware {

    @Value("${xddcloud.gateway.route.sentinel.qpscount:10}")
    private int qpscount;

    @Value("${xddcloud.gateway.route.sentinel.interval:1}")
    private int interval;


    private static final String GATEWAY_ROUTE_PREFIX = "ReactiveCompositeDiscoveryClient_";


    private final List<ViewResolver> viewResolvers;

    private final ServerCodecConfigurer serverCodecConfigurer;

    public GatewaySentinelConfiguration(ObjectProvider<List<ViewResolver>> viewResolversProvider,
        ServerCodecConfigurer serverCodecConfigurer) {
        this.viewResolvers = viewResolversProvider.getIfAvailable(Collections::emptyList);
        this.serverCodecConfigurer = serverCodecConfigurer;
    }

    @Bean
    @Order(Ordered.HIGHEST_PRECEDENCE)
    public SentinelGatewayBlockExceptionHandler sentinelGatewayBlockExceptionHandler() {
        return new SentinelGatewayBlockExceptionHandler(viewResolvers, serverCodecConfigurer);
    }

    @Bean
    @Order(Ordered.HIGHEST_PRECEDENCE)
    public SentinelGatewayFilter sentinelGatewayFilter() {
        log.info("[Sentinel SpringCloudGateway] register SentinelGatewayFilter with order: {}",
            Ordered.HIGHEST_PRECEDENCE);
        return new SentinelGatewayFilter(Ordered.HIGHEST_PRECEDENCE);
    }


    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        // 初始化流控规则
        Set<GatewayFlowRule> rules = new HashSet<>();
        Flux<String> services = applicationContext.getBean(ReactiveDiscoveryClient.class).getServices();
        Iterable<GatewayFlowRule> iterable = services.flatMap(service -> Flux.just(
            new GatewayFlowRule(GATEWAY_ROUTE_PREFIX.concat(service)).setGrade(RuleConstant.FLOW_GRADE_QPS)
                .setIntervalSec(interval) // 统计时间窗口，单位是秒，默认是 1 秒
                .setCount(qpscount))// 限流阈值;
        ).toIterable();
        iterable.forEach(rules::add);
        GatewayRuleManager.loadRules(rules);

        // 初始化拦截处理器
        BlockRequestHandler blockRequestHandler = (serverWebExchange, throwable) -> {
            CommonResult<String> commonResult = CommonResult.failed("too many request");
            commonResult.setData(serverWebExchange.getRequest().getPath().toString());
            commonResult.setCode(HttpStatus.TOO_MANY_REQUESTS.value());
            return ServerResponse.status(HttpStatus.TOO_MANY_REQUESTS).contentType(MediaType.APPLICATION_JSON_UTF8)
                .body(BodyInserters.fromValue(commonResult));
        };
        GatewayCallbackManager.setBlockHandler(blockRequestHandler);
    }
}
